import os
from time import sleep

try:
    from anthropic import Anthropic
except ImportError as e:
    pass

from lcb_runner.runner.base_runner import BaseRunner


class Claude3Runner(BaseRunner):
    client = Anthropic(api_key=os.getenv("ANTHROPIC_KEY"), timeout=1200)

    def __init__(self, args, model):
        super().__init__(args, model)
        if "Thinking" in model.model_style.value:
            self.client_kwargs: dict[str | str] = {
                "model": args.model,
                "max_tokens": 32000,
                "thinking": {"type": "enabled", "budget_tokens": 24000},
                "stream": False,
            }
        else:
            self.client_kwargs: dict[str | str] = {
                "model": args.model,
                "temperature": args.temperature,
                "max_tokens": args.max_tokens,
                "top_p": args.top_p,
            }

    def _run_single(self, prompt: tuple[str, str]) -> list[str]:

        def __run_single(counter):
            try:
                response = self.client.messages.create(
                    system=prompt[0],
                    messages=prompt[1],
                    **self.client_kwargs,
                )
                content = "\n".join(
                    [
                        getattr(x, "text", getattr(x, "thinking", "\nREDACTED\n"))
                        for x in response.content
                    ]
                )
                return content
            except Exception as e:
                print("Exception: ", repr(e), "Sleeping for 20 seconds...")
                sleep(20 * (11 - counter))
                counter = counter - 1
                if counter == 0:
                    print(f"Failed to run model for {prompt}!")
                    print("Exception: ", repr(e))
                    raise e
                return __run_single(counter)

        outputs = []
        try:
            for _ in range(self.args.n):
                outputs.append(__run_single(10))
        except Exception as e:
            raise e

        return outputs
